import argparse
from functools import cached_property
import logging
from pathlib import Path
from typing import Any, Generator
from pydantic import BaseModel

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("Split FASTA")


class SplitFasta(BaseModel):
    baseG: float
    fasta_file: Path
    out_prefix: str

    @cached_property
    def base_num(self) -> int:
        return int(self.baseG * 1e9)

    def iter_fasta(self) -> Generator[tuple[str, str, int], Any, None]:
        with open(self.fasta_file) as reader:
            acc, seq, num = "", "", 0
            for line in reader:
                if line[0] == ">":
                    if acc:
                        yield acc, seq, num
                    acc, seq, num = line[1:].split(" ", 1)[0], "", 0
                else:
                    seq += line
                    num += len(line) - 1
            yield acc, seq, num

    def run(self) -> Path:
        Path(self.out_prefix).parent.mkdir(parents=True, exist_ok=True)
        index, base_num = 1, 0
        writer = open(f"{self.out_prefix}.000.fa", "w")
        logging.info("Write %s.000.fa", self.out_prefix)
        for acc, seq, num in self.iter_fasta():
            writer.write(f">{acc}\n{seq}")
            base_num += num
            if base_num >= self.base_num:
                writer.close()
                writer = open(f"{self.out_prefix}.{index:03d}.fa", "w")
                logging.info("Write %s.%s.fa", self.out_prefix, f"{index:03d}")
                index += 1
                base_num = 0
        writer.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="gcGenome")
    parser.add_argument("-i", "--fasta_file", type=Path, required=True, help="input fasta file")
    parser.add_argument("-o", "--out_prefix", type=str, required=True, help="output prefix")
    parser.add_argument("-b", "--baseG", type=float, default=3, help="each base size")
    parser.set_defaults(func=lambda args: SplitFasta.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)
